import os

import SimpleITK as sitk
from PIL import Image
import scipy.ndimage as ndimage
import numpy as np
'''

BTCV的标签
背景：0
肝（liver）：6
右肾（right kidney）：2
左肾（left kidney）：3
脾（spleen）：1
'''
def processAndSave(data_dir=r"D:\AllData\BTCV\RawData\Training",
                   save_dir=r"D:\AllData\BTCV\afterProcess"):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    number = 0
    label_dir = os.path.join(data_dir,"volume" )
    image_dir = os.path.join(data_dir,"segmentation" )
    subject_list=os.listdir(image_dir)
    subject_list.sort()
    for index,nii_name in enumerate(subject_list):
        # -------------------处理图像数据---------------------#
        image_path=os.path.join(image_dir,nii_name)
        image_nii=sitk.ReadImage(image_path)
        image_array=sitk.GetArrayFromImage(image_nii)
        last_shape=np.array([len(image_array),256,256])
        image_dimension_adjustment = 1 / (image_array.shape / last_shape)
        image_array = ndimage.zoom(image_array, image_dimension_adjustment, order=3)  # 双线性插值
        image_array = (image_array - image_array.mean()) / image_array.std()

        # -------------------处理标注数据---------------------#
        label_path=os.path.join(label_dir,nii_name).replace('volume','segmentation')
        label_nii = sitk.ReadImage(label_path)
        label_array = sitk.GetArrayFromImage(label_nii)
        label_array=label_array.astype('uint8')
        # 将脾的标签设为1，其余器官和背景设为0
        label_array[label_array != 1] = 0
        label_array[label_array == 1] = 1

        label_dimension_adjustment = 1 / (label_array.shape / last_shape)

        label_array = ndimage.zoom(label_array, label_dimension_adjustment, order=0)  # 只会出现原来出现过的数值

        #--------拼接image和label的每个2D切片，[image, mask] 。并保存-----#
        for i in range(len(image_array)):
            if 1 in label_array[i]:
                number += 1
                # 拼接并保存
                # combine_array = np.array([image_array[i], label_array[i]])
                n="00"+str(index+1) if index<9 else "0"+str(index+1)
                last_dir=os.path.join(save_dir,"s"+n)
                if not os.path.exists(last_dir):
                    os.makedirs(last_dir)
                save_name = os.path.join(last_dir, str(i+1) + '.npz')
                np.savez(save_name, arr_0=image_array[i],arr_1=label_array[i])#为了跟代码的dateset对应，参数名和dataset统一arr_0，arr_1
        print(index+1)
    print(number)
if __name__ == '__main__':
    #gpu2
    processAndSave(data_dir=r"/home/liukai/AllData/zhongShanBreast/nii",
                   save_dir=r"/home/liukai/AllData/BreastDomain/ForSIFA/subjects/zhongShanBreast")

    processAndSave(data_dir=r"/home/liukai/AllData/zunYiBreast/nii",
                   save_dir=r"/home/liukai/AllData/BreastDomain/ForSIFA/subjects/zunYi")

    # # gpu3
    # processAndSave(data_dir=r"/home/liukai/AllData/BTCV/RawData/Training",
    #                save_dir=r"/home/liukai/AllData/AbdominalOrgansForDomain/ForSIFA/BTCV")

